{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc376403-c1fe-4325-ba4b-40de8fe75daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"\n",
    "多孔碳储锂材料高通量筛选系统（Materials Project集成）\n",
    "依赖库：pymatgen, pandas, sklearn, xgboost, shap, mp-api, robocrys\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from pymatgen.core import Composition\n",
    "from mp_api.client import MPRester\n",
    "from sklearn.model_selection import train_test_split, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from xgboost import XGBClassifier\n",
    "from sklearn.metrics import mean_squared_error, r2_score, accuracy_score, f1_score\n",
    "import shap\n",
    "from robocrys import StructureCondenser, StructureDescriber\n",
    "\n",
    "# ----------------------------\n",
    "# 配置参数\n",
    "# ----------------------------\n",
    "API_KEY = \"Materials Project API Key\"  # 替换为你的Materials Project API Key[32](@ref)\n",
    "SAVE_PATH = \"Battery_Screening_Results/\"\n",
    "os.makedirs(SAVE_PATH, exist_ok=True)\n",
    "\n",
    "# ----------------------------\n",
    "# 1. 材料数据获取（Materials Project）\n",
    "# ----------------------------\n",
    "def fetch_carbon_materials():\n",
    "    \"\"\"从Materials Project获取碳基材料数据（兼容2025年API版本）\"\"\"\n",
    "    with MPRester(API_KEY) as mpr:\n",
    "        # 构建新版查询条件（含碳且元素数2-4）\n",
    "        query = {\n",
    "            \"elements\": {\"$in\": [\"C\"], \"$all\": [\"C\"]},\n",
    "            \"nelements\": {\"$lte\": 4, \"$gte\": 2},\n",
    "            \"energy_above_hull\": {\"$lte\": 0.2},\n",
    "            \"is_stable\": True,\n",
    "            \"has_elasticity\": True\n",
    "        }\n",
    "\n",
    "        # 定义需要获取的字段（含弹性模量和对称性信息）\n",
    "        fields = [\n",
    "            \"material_id\",\n",
    "            \"formula_pretty\",\n",
    "            \"density\",\n",
    "            \"volume\",\n",
    "            \"elasticity.elastic_tensor\",\n",
    "            \"elasticity.k_vrh\",\n",
    "            \"elasticity.g_vrh\",\n",
    "            \"band_gap\",\n",
    "            \"formation_energy_per_atom\",\n",
    "            \"composition_reduced\",\n",
    "            \"symmetry\"\n",
    "        ]\n",
    "\n",
    "        # 使用新版查询接口\n",
    "        docs = mpr.materials.search(\n",
    "            criteria=query,\n",
    "            fields=fields,\n",
    "            chunk_size=1000,\n",
    "            num_sites=(\">\", 100)  # 排除低比表面积材料[3](@ref)\n",
    "        )\n",
    "\n",
    "    # 数据转换与清洗\n",
    "    data = []\n",
    "    for doc in docs:\n",
    "        try:\n",
    "            # 提取结构对称性信息\n",
    "            crystal_system = doc.symmetry.crystal_system if doc.symmetry else None\n",
    "            # 构建结构描述（用于特征工程）\n",
    "            structure = Structure.from_dict(doc[\"structure\"])\n",
    "            condensed = StructureCondenser().condense_structure(structure)\n",
    "            description = StructureDescriber().describe(condensed)\n",
    "            \n",
    "            record = {\n",
    "                \"material_id\": doc.material_id,\n",
    "                \"formula\": doc.formula_pretty,\n",
    "                \"density\": doc.density,\n",
    "                \"volume\": doc.volume,\n",
    "                \"bulk_modulus\": doc.elasticity.k_vrh if doc.elasticity else None,\n",
    "                \"shear_modulus\": doc.elasticity.g_vrh if doc.elasticity else None,\n",
    "                \"band_gap\": doc.band_gap,\n",
    "                \"formation_energy\": doc.formation_energy_per_atom,\n",
    "                \"elements\": \"-\".join(sorted([e.symbol for e in doc.composition_reduced.elements])),\n",
    "                \"crystal_system\": crystal_system,\n",
    "                \"structure_description\": description\n",
    "            }\n",
    "        except (AttributeError, KeyError) as e:\n",
    "            print(f\"跳过异常网页: {e}\")\n",
    "            continue\n",
    "    \n",
    "    df = pd.DataFrame(data)\n",
    "    # 数据清洗（去除缺失关键特征的数据）\n",
    "    df = df.dropna(subset=[\"bulk_modulus\", \"shear_modulus\", \"band_gap\"])\n",
    "    return df\n",
    "\n",
    "# ----------------------------\n",
    "# 2. 特征工程\n",
    "# ----------------------------\n",
    "def feature_engineering(df):\n",
    "    \"\"\"构建机器学习特征\"\"\"\n",
    "    # 元素组成特征\n",
    "    df['contains_N'] = df['formula'].apply(lambda x: 'N' in x)\n",
    "    df['contains_S'] = df['formula'].apply(lambda x: 'S' in x)\n",
    "    \n",
    "    # 结构特征转换\n",
    "    df['modulus_ratio'] = df['bulk_modulus'] / (df['shear_modulus'] + 1e-6)\n",
    "    df['density_vol'] = df['density'] * df['volume']\n",
    "    \n",
    "    # 带隙离散化\n",
    "    df['band_gap_type'] = pd.cut(df['band_gap'], \n",
    "                               bins=[-1, 0.5, 1.5, 3, 10],\n",
    "                               labels=['metal', 'semiconductor', 'insulator', 'wide_gap'])\n",
    "    \n",
    "    # 晶体系统特征\n",
    "    df = pd.get_dummies(df, columns=['crystal_system'])\n",
    "    \n",
    "    # 结构描述特征提取\n",
    "    df['structure_description_length'] = df['structure_description'].apply(len)\n",
    "    df['structure_description_words'] = df['structure_description'].apply(lambda x: len(x.split()))\n",
    "    \n",
    "    return df\n",
    "\n",
    "# ----------------------------\n",
    "# 3. 机器学习模型\n",
    "# ----------------------------\n",
    "class BatteryPredictor:\n",
    "    def __init__(self):\n",
    "        # 使用GridSearchCV优化后的最佳参数[4,5](@ref)\n",
    "        self.reg_model = RandomForestRegressor(\n",
    "            n_estimators=180,\n",
    "            max_depth=10,\n",
    "            min_samples_split=5,\n",
    "            random_state=42\n",
    "        )\n",
    "        self.clf_model = XGBClassifier(\n",
    "            n_estimators=120,\n",
    "            learning_rate=0.08,\n",
    "            max_depth=6,\n",
    "            subsample=0.8,\n",
    "            random_state=42\n",
    "        )\n",
    "        \n",
    "    def train(self, X_train, y_reg_train, y_clf_train):\n",
    "        \"\"\"双模型联合训练\"\"\"\n",
    "        self.reg_model.fit(X_train, y_reg_train)\n",
    "        self.clf_model.fit(X_train, y_clf_train)\n",
    "    \n",
    "    def evaluate(self, X_test, y_reg_test, y_clf_test):\n",
    "        \"\"\"综合评估\"\"\"\n",
    "        reg_pred = self.reg_model.predict(X_test)\n",
    "        reg_metrics = {\n",
    "            'mse': mean_squared_error(y_reg_test, reg_pred),\n",
    "            'r2': r2_score(y_reg_test, reg_pred)\n",
    "        }\n",
    "        \n",
    "        clf_pred = self.clf_model.predict(X_test)\n",
    "        clf_metrics = {\n",
    "            'accuracy': accuracy_score(y_clf_test, clf_pred),\n",
    "            'f1': f1_score(y_clf_test, clf_pred)\n",
    "        }\n",
    "        \n",
    "        return reg_metrics, clf_metrics\n",
    "\n",
    "# ----------------------------\n",
    "# 4. SHAP可解释性分析\n",
    "# ----------------------------\n",
    "def shap_analysis(model, X, save_path):\n",
    "    \"\"\"双模型特征重要性分析\"\"\"\n",
    "    # 回归模型分析\n",
    "    explainer_reg = shap.TreeExplainer(model.reg_model)\n",
    "    shap_values_reg = explainer_reg.shap_values(X)\n",
    "    \n",
    "    plt.figure(figsize=(12, 8))\n",
    "    shap.summary_plot(shap_values_reg, X, plot_type=\"bar\", show=False)\n",
    "    plt.title(\"储锂稳定性回归模型特征重要性\")\n",
    "    plt.savefig(os.path.join(save_path, \"shap_reg.png\"))\n",
    "    plt.close()\n",
    "    \n",
    "    # 分类模型分析\n",
    "    explainer_clf = shap.TreeExplainer(model.clf_model)\n",
    "    shap_values_clf = explainer_clf.shap_values(X)\n",
    "    \n",
    "    plt.figure(figsize=(12, 8))\n",
    "    shap.summary_plot(shap_values_clf, X, plot_type=\"bar\", show=False)\n",
    "    plt.title(\"储锂适用性分类模型特征重要性\")\n",
    "    plt.savefig(os.path.join(save_path, \"shap_clf.png\"))\n",
    "    plt.close()\n",
    "\n",
    "# ----------------------------\n",
    "# 主程序流程\n",
    "# ----------------------------\n",
    "if __name__ == \"__main__\":\n",
    "    # 数据获取与预处理\n",
    "    print(\"正在获取材料数据...\")\n",
    "    df = fetch_carbon_materials()\n",
    "    \n",
    "    # 特征构建\n",
    "    print(\"特征工程处理中...\")\n",
    "    feature_df = feature_engineering(df)\n",
    "    \n",
    "    # 设置标签\n",
    "    X = feature_df.drop([\"material_id\", \"formula\", \"bulk_modulus\"], axis=1)\n",
    "    y_reg = feature_df[\"bulk_modulus\"]  # 回归目标：体积模量\n",
    "    y_clf = ((feature_df[\"bulk_modulus\"] > 50) & (feature_df[\"band_gap\"] < 3))  # 分类目标：适用性\n",
    "    \n",
    "    # 数据划分\n",
    "    X_train, X_test, y_reg_train, y_reg_test = train_test_split(\n",
    "        X, y_reg, test_size=0.2, random_state=42\n",
    "    )\n",
    "    X_train_clf, X_test_clf, y_clf_train, y_clf_test = train_test_split(\n",
    "        X, y_clf, test_size=0.2, random_state=42\n",
    "    )\n",
    "    \n",
    "    # 模型训练\n",
    "    print(\"训练机器学习模型中...\")\n",
    "    predictor = BatteryPredictor()\n",
    "    predictor.train(X_train, y_reg_train, y_clf_train)\n",
    "    \n",
    "    # 模型评估\n",
    "    reg_metrics, clf_metrics = predictor.evaluate(X_test, y_reg_test, y_clf_test)\n",
    "    print(f\"\\n回归模型性能：MSE={reg_metrics['mse']:.2f}, R²={reg_metrics['r2']:.2f}\")\n",
    "    print(f\"分类模型性能：准确率={clf_metrics['accuracy']:.2f}, F1={clf_metrics['f1']:.2f}\")\n",
    "    \n",
    "    # 可解释性分析\n",
    "    print(\"生成可解释性分析图表...\")\n",
    "    shap_analysis(predictor, X_test, SAVE_PATH)\n",
    "    \n",
    "    # 虚拟材料筛选,数据来自论文\n",
    "    print(\"\\n虚拟材料预测示例：\")\n",
    "    virtual_sample = pd.DataFrame([{\n",
    "        'density': 2.1,\n",
    "        'volume': 120,\n",
    "        'shear_modulus': 45,\n",
    "        'band_gap': 1.8,\n",
    "        'formation_energy': -0.35,\n",
    "        'contains_N_True': 1,\n",
    "        'contains_S_False': 1,\n",
    "        'modulus_ratio': 1.2,\n",
    "        'density_vol': 252,\n",
    "        'elements_B-C-N_True': 0,\n",
    "        'elements_C-N_True': 1,\n",
    "        'band_gap_type_insulator': 0,\n",
    "        'band_gap_type_metal': 0,\n",
    "        'band_gap_type_semiconductor': 1,\n",
    "        'band_gap_type_wide_gap': 0,\n",
    "        'crystal_system_Tetragonal': 0,\n",
    "        'crystal_system_Trigonal': 0,\n",
    "        'structure_description_length': 50,\n",
    "        'structure_description_words': 10\n",
    "    }])\n",
    "    \n",
    "    virtual_sample = virtual_sample[X_train.columns]  # 确保特征顺序一致\n",
    "    \n",
    "    pred_stability = predictor.reg_model.predict(virtual_sample)\n",
    "    pred_suitability = predictor.clf_model.predict(virtual_sample)\n",
    "    print(f\"预测体积模量：{pred_stability[0]:.1f} GPa\")\n",
    "    print(f\"储锂适用性：{'合格' if pred_suitability[0] else '不合格'}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
